from typing import Iterable

from pyflink.common import Types, Time, SimpleStringSchema, WatermarkStrategy
from pyflink.datastream import StreamExecutionEnvironment, RuntimeExecutionMode, MapFunction, ProcessFunction, \
    KeyedProcessFunction, ProcessWindowFunction
from pyflink.datastream.connectors.kafka import KafkaSource, KafkaOffsetsInitializer
from pyflink.datastream.functions import KEY, IN, OUT
from pyflink.datastream.window import TumblingProcessingTimeWindows

# 1、创建flink执行环境
env = StreamExecutionEnvironment.get_execution_environment()

# 修改并行度
env.set_parallelism(1)

source = KafkaSource.builder() \
    .set_bootstrap_servers("master:9092") \
    .set_topics("students_csv") \
    .set_group_id("my-group") \
    .set_starting_offsets(KafkaOffsetsInitializer.earliest()) \
    .set_value_only_deserializer(SimpleStringSchema()) \
    .build()

# 基于kafka构建DataStream (无界流)
lines_ds = env.from_source(source, WatermarkStrategy.no_watermarks(), "Kafka Source")


# process算子，底层算子，可以操作底层的状态，事件，时间
# 1、在普通的DS上使用process, 可以代替map flatMap filter功能

class MapProcessFunction(ProcessFunction):

    # process_element： 每一条数据执行一次，
    # line：一行数据
    # ctx： 上下文对象
    def process_element(self, line, ctx: 'ProcessFunction.Context'):
        # 在process中可以获取数据的时间
        ts = ctx.timestamp()
        print(f"数据时间：{ts}")
        # 使用yield返回，一次返回一行
        yield line.split(",")[-1]


clazz_ds = lines_ds.process(MapProcessFunction(), output_type=Types.STRING())

key_by_ds = clazz_ds.key_by(lambda clazz: clazz)


# 2、在分组之后使用，数据已经被分到同一个task中了，可以在process进行聚合计算

class CountProcessFunction(KeyedProcessFunction):

    def __init__(self):
        # count属性，每一个task中所有的key共用同一个
        # self.count = 0

        # 定义字典，为每一个班级都保存一个数量
        self.count_dict = {}

    # process_element： 每一条数据执行一次，
    # line：一行数据
    # ctx： 上下文对象
    def process_element(self, line, ctx: 'KeyedProcessFunction.Context'):
        # 获取分组的key
        clazz = ctx.get_current_key()

        if clazz not in self.count_dict:
            self.count_dict[clazz] = 1
        else:
            self.count_dict[clazz] += 1

        # 返回统计结果
        yield clazz, self.count_dict[clazz]


clazz_num_ds = key_by_ds.process(CountProcessFunction())

# clazz_num_ds.print()

# 3、在窗口之后使用
window_ds = key_by_ds.window(TumblingProcessingTimeWindows.of(Time.seconds(5)))


class CountProcessWindowFunction(ProcessWindowFunction):

    # clazz: 分组的key
    # context: 上下文对象
    # elements： 窗口内所有的数据
    def process(self, clazz: KEY, context: 'ProcessWindowFunction.Context', elements: Iterable[IN]) -> Iterable[OUT]:
        window = context.window()

        # 获取窗口的开始和结束时间
        start = window.start
        end = window.end

        # 统计
        num = 0
        for line in elements:
            num += 1

        yield clazz, start, end, num


window_ds.process(CountProcessWindowFunction()).print()

env.execute()
